#include "md5_8.h"
#include <iomanip>
#include <assert.h>
#include <chrono>
#include <algorithm>
#include <arm_sve.h>

using namespace std;
using namespace chrono;

// 重用md5.cpp中的StringProcess函数
Byte *StringProcess(string input, int *n_byte)
{
    // 将输入的字符串转换为Byte为单位的数组
    Byte *blocks = (Byte *)input.c_str();
    int length = input.length();

    // 计算原始消息长度（以比特为单位）
    int bitLength = length * 8;

    // paddingBits: 原始消息需要的padding长度（以bit为单位）
    int paddingBits = bitLength % 512;
    if (paddingBits > 448)
    {
        paddingBits += 512 - (paddingBits - 448);
    }
    else if (paddingBits < 448)
    {
        paddingBits = 448 - paddingBits;
    }
    else if (paddingBits == 448)
    {
        paddingBits = 512;
    }

    // 原始消息需要的padding长度（以Byte为单位）
    int paddingBytes = paddingBits / 8;
    // 创建最终的字节数组
    int paddedLength = length + paddingBytes + 8;
    Byte *paddedMessage = new Byte[paddedLength];

    // 复制原始消息
    memcpy(paddedMessage, blocks, length);

    // 添加填充字节
    paddedMessage[length] = 0x80;
    memset(paddedMessage + length + 1, 0, paddingBytes - 1);

    // 添加消息长度（64比特，小端格式）
    for (int i = 0; i < 8; ++i)
    {
        paddedMessage[length + paddingBytes + i] = ((uint64_t)length * 8 >> (i * 8)) & 0xFF;
    }

    *n_byte = paddedLength;
    return paddedMessage;
}

void SIMD_MD5Hash_8(string input[], bit32 state[8][4])
{
    Byte *paddedMessage[8];
    int *messageLength = new int[8];
    
    
    // 处理所有8个输入字符串
    for (int i = 0; i < 8; i++)
    {
        paddedMessage[i] = StringProcess(input[i], &messageLength[i]);
    }
    
    // 计算每个输入需要的块数
    int n_blocks[8];
    for (int i = 0; i < 8; i++) {
        n_blocks[i] = messageLength[i] / 64;
    }
    
    // 找出最大和最小块数
    int max_n_blocks = n_blocks[0];
    int min_n_blocks = n_blocks[0];
    for (int i = 1; i < 8; i++) {
        if (n_blocks[i] > max_n_blocks) max_n_blocks = n_blocks[i];
        if (n_blocks[i] < min_n_blocks) min_n_blocks = n_blocks[i];
        std::cout<<"n_blocks["<<i<<"] = "<<n_blocks[i]<<std::endl;
    }

    // 初始化所有8个状态
    for (int i = 0; i < 8; i++)
    {
        state[i][0] = 0x67452301;
        state[i][1] = 0xefcdab89;
        state[i][2] = 0x98badcfe;
        state[i][3] = 0x10325476;
    }
  
    // SVE向量不能作为数组元素，需要分开声明
    svbool_t pg = svptrue_b32(); // 使用这个而不是 svwhilelt_b32(0, 8)
    
    // // 初始化SVE向量
    // svuint32_t a_end = svdup_n_u32_z(pg, 0x67452301);
    // svuint32_t b_end = svdup_n_u32_z(pg, 0xefcdab89);
    // svuint32_t c_end = svdup_n_u32_z(pg, 0x98badcfe);
    // svuint32_t d_end = svdup_n_u32_z(pg, 0x10325476);

    // 在第一轮循环前，重新初始化SVE向量确保8个Lane全部正确
    uint32_t init_a[8] = {0x67452301, 0x67452301, 0x67452301, 0x67452301, 0x67452301, 0x67452301, 0x67452301, 0x67452301};
    uint32_t init_b[8] = {0xefcdab89, 0xefcdab89, 0xefcdab89, 0xefcdab89, 0xefcdab89, 0xefcdab89, 0xefcdab89, 0xefcdab89};
    uint32_t init_c[8] = {0x98badcfe, 0x98badcfe, 0x98badcfe, 0x98badcfe, 0x98badcfe, 0x98badcfe, 0x98badcfe, 0x98badcfe};
    uint32_t init_d[8] = {0x10325476, 0x10325476, 0x10325476, 0x10325476, 0x10325476, 0x10325476, 0x10325476, 0x10325476};

    svuint32_t a_end = svld1_u32(pg, init_a);
    svuint32_t b_end = svld1_u32(pg, init_b);
    svuint32_t c_end = svld1_u32(pg, init_c);
    svuint32_t d_end = svld1_u32(pg, init_d);

    svuint32_t a = a_end;
    svuint32_t b = b_end;
    svuint32_t c = c_end; 
    svuint32_t d = d_end;



    // 不能将SVE类型作为数组元素，所以必须单独声明每个向量
    svuint32_t x0, x1, x2, x3, x4, x5, x6, x7;
    svuint32_t x8, x9, x10, x11, x12, x13, x14, x15;
    
    // 跟踪已完成处理的输入数量
    int over_num = 0;
    bool finished[8] = {false, false, false, false, false, false, false, false};
    
    // 逐块处理
    for (int i = 0; i <= max_n_blocks; i++)
    {
        // 或者在每次循环前验证 pg 是否有效
        int pg_lanes = svcntp_b32(svptrue_b32(), pg);
        std::cout << "Active lanes in pg: " << pg_lanes << std::endl;

        // 检查每个输入是否已经处理完成
        for (int j = 0; j < 8&&(i+1)>=min_n_blocks; j++) {
            if (!finished[j] && i == n_blocks[j]) {
                // 从SVE向量中提取结果 - 使用lane访问
                uint32_t tmp[4][8];
                svst1_u32(pg, tmp[0], a_end);
                svst1_u32(pg, tmp[1], b_end);
                svst1_u32(pg, tmp[2], c_end);
                svst1_u32(pg, tmp[3], d_end);
                state[j][0] = tmp[0][j];
                state[j][1] = tmp[1][j];
                state[j][2] = tmp[2][j];
                state[j][3] = tmp[3][j];
                // state[j][0] = svget_lane_u32(a_end, j);
                // state[j][1] = svget_lane_u32(b_end, j);
                // state[j][2] = svget_lane_u32(c_end, j);
                // state[j][3] = svget_lane_u32(d_end, j);
                
                finished[j] = true;
                over_num++;
            }
        }
        

            if (over_num == 8)
        {
          break; // 所有输入都处理完毕
        }

        // 为每个16个字处理块
        uint32_t temp_x[16][8] = {0};
        
        // 准备所有输入的所有16个字
        for (int i1 = 0; i1 < 16; ++i1)
        {
            int offset = 4 * i1 + i * 64;
            
            for (int j = 0; j < 8; j++) {
                if (i < n_blocks[j]) {
                    //只处理未完成的输入
                    temp_x[i1][j] = *reinterpret_cast<uint32_t *>(paddedMessage[j] + offset);
                }
            }
            std::cout<<temp_x[i1][0]<<" "<<temp_x[i1][1]<<" "<<temp_x[i1][2]<<" "<<temp_x[i1][3]<<" "<<temp_x[i1][4]<<" "<<temp_x[i1][5]<<" "<<temp_x[i1][6]<<" "<<temp_x[i1][7]<<std::endl;
        }
        
        // 分别加载每个向量
        x0 = svld1_u32(pg, temp_x[0]);
        x1 = svld1_u32(pg, temp_x[1]);
        x2 = svld1_u32(pg, temp_x[2]);
        x3 = svld1_u32(pg, temp_x[3]);
        x4 = svld1_u32(pg, temp_x[4]);
        x5 = svld1_u32(pg, temp_x[5]);
        x6 = svld1_u32(pg, temp_x[6]);
        x7 = svld1_u32(pg, temp_x[7]);
        x8 = svld1_u32(pg, temp_x[8]);
        x9 = svld1_u32(pg, temp_x[9]);
        x10 = svld1_u32(pg, temp_x[10]);
        x11 = svld1_u32(pg, temp_x[11]);
        x12 = svld1_u32(pg, temp_x[12]);
        x13 = svld1_u32(pg, temp_x[13]);
        x14 = svld1_u32(pg, temp_x[14]);
        x15 = svld1_u32(pg, temp_x[15]);

        // 使用具体变量替代数组访问
        /* Round 1 */
        FF_SVE(a, b, c, d, x0, s11, 0xd76aa478, pg);
        FF_SVE(d, a, b, c, x1, s12, 0xe8c7b756, pg);
        FF_SVE(c, d, a, b, x2, s13, 0x242070db, pg);
        FF_SVE(b, c, d, a, x3, s14, 0xc1bdceee, pg);
        FF_SVE(a, b, c, d, x4, s11, 0xf57c0faf, pg);
        FF_SVE(d, a, b, c, x5, s12, 0x4787c62a, pg);
        FF_SVE(c, d, a, b, x6, s13, 0xa8304613, pg);
        FF_SVE(b, c, d, a, x7, s14, 0xfd469501, pg);
        FF_SVE(a, b, c, d, x8, s11, 0x698098d8, pg);
        FF_SVE(d, a, b, c, x9, s12, 0x8b44f7af, pg);
        FF_SVE(c, d, a, b, x10, s13, 0xffff5bb1, pg);
        FF_SVE(b, c, d, a, x11, s14, 0x895cd7be, pg);
        FF_SVE(a, b, c, d, x12, s11, 0x6b901122, pg);
        FF_SVE(d, a, b, c, x13, s12, 0xfd987193, pg);
        FF_SVE(c, d, a, b, x14, s13, 0xa679438e, pg);
        FF_SVE(b, c, d, a, x15, s14, 0x49b40821, pg);

        /* Round 2 */
        GG_SVE(a, b, c, d, x1, s21, 0xf61e2562, pg);
        GG_SVE(d, a, b, c, x6, s22, 0xc040b340, pg);
        GG_SVE(c, d, a, b, x11, s23, 0x265e5a51, pg);
        GG_SVE(b, c, d, a, x0, s24, 0xe9b6c7aa, pg);
        GG_SVE(a, b, c, d, x5, s21, 0xd62f105d, pg);
        GG_SVE(d, a, b, c, x10, s22, 0x2441453, pg);
        GG_SVE(c, d, a, b, x15, s23, 0xd8a1e681, pg);
        GG_SVE(b, c, d, a, x4, s24, 0xe7d3fbc8, pg);
        GG_SVE(a, b, c, d, x9, s21, 0x21e1cde6, pg);
        GG_SVE(d, a, b, c, x14, s22, 0xc33707d6, pg);
        GG_SVE(c, d, a, b, x3, s23, 0xf4d50d87, pg);
        GG_SVE(b, c, d, a, x8, s24, 0x455a14ed, pg);
        GG_SVE(a, b, c, d, x13, s21, 0xa9e3e905, pg);
        GG_SVE(d, a, b, c, x2, s22, 0xfcefa3f8, pg);
        GG_SVE(c, d, a, b, x7, s23, 0x676f02d9, pg);
        GG_SVE(b, c, d, a, x12, s24, 0x8d2a4c8a, pg);
            /* Round 3 */
        HH_SVE(a, b, c, d, x5, s31, 0xfffa3942, pg);
        HH_SVE(d, a, b, c, x8, s32, 0x8771f681, pg);
        HH_SVE(c, d, a, b, x11, s33, 0x6d9d6122, pg);
        HH_SVE(b, c, d, a, x14, s34, 0xfde5380c, pg);
        HH_SVE(a, b, c, d, x1, s31, 0xa4beea44, pg);
        HH_SVE(d, a, b, c, x4, s32, 0x4bdecfa9, pg);
        HH_SVE(c, d, a, b, x7, s33, 0xf6bb4b60, pg);
        HH_SVE(b, c, d, a, x10, s34, 0xbebfbc70, pg);
        HH_SVE(a, b, c, d, x13, s31, 0x289b7ec6, pg);
        HH_SVE(d, a, b, c, x0, s32, 0xeaa127fa, pg);
        HH_SVE(c, d, a, b, x3, s33, 0xd4ef3085, pg);
        HH_SVE(b, c, d, a, x6, s34, 0x4881d05, pg);
        HH_SVE(a, b, c, d, x9, s31, 0xd9d4d039, pg);
        HH_SVE(d, a, b, c, x12, s32, 0xe6db99e5, pg);
        HH_SVE(c, d, a, b, x15, s33, 0x1fa27cf8, pg);
        HH_SVE(b, c, d, a, x2, s34, 0xc4ac5665, pg);

        /* Round 4 */
        II_SVE(a, b, c, d, x0, s41, 0xf4292244, pg);
        II_SVE(d, a, b, c, x7, s42, 0x432aff97, pg);
        II_SVE(c, d, a, b, x14, s43, 0xab9423a7, pg);
        II_SVE(b, c, d, a, x5, s44, 0xfc93a039, pg);
        II_SVE(a, b, c, d, x12, s41, 0x655b59c3, pg);
        II_SVE(d, a, b, c, x3, s42, 0x8f0ccc92, pg);
        II_SVE(c, d, a, b, x10, s43, 0xffeff47d, pg);
        II_SVE(b, c, d, a, x1, s44, 0x85845dd1, pg);
        II_SVE(a, b, c, d, x8, s41, 0x6fa87e4f, pg);
        II_SVE(d, a, b, c, x15, s42, 0xfe2ce6e0, pg);
        II_SVE(c, d, a, b, x6, s43, 0xa3014314, pg);
        II_SVE(b, c, d, a, x13, s44, 0x4e0811a1, pg);
        II_SVE(a, b, c, d, x4, s41, 0xf7537e82, pg);
        II_SVE(d, a, b, c, x11, s42, 0xbd3af235, pg);
        II_SVE(c, d, a, b, x2, s43, 0x2ad7d2bb, pg);
        II_SVE(b, c, d, a, x9, s44, 0xeb86d391, pg);

        // 更新状态
        a_end = svadd_z(pg, a_end, a);
        b_end = svadd_z(pg, b_end, b);
        c_end = svadd_z(pg, c_end, c);
        d_end = svadd_z(pg, d_end, d);
        
        a = a_end;
        b = b_end;
        c = c_end;
        d = d_end;

        // 打印每次循环后的结果
        uint32_t tmp[4][8];
        svst1_u32(pg, tmp[0], a_end);
        svst1_u32(pg, tmp[1], b_end);
        svst1_u32(pg, tmp[2], c_end);
        svst1_u32(pg, tmp[3], d_end);
        cout << "After loop iteration " << i << ":" << endl;
        for (int j = 0; j < 8; j++)
        {
          cout << "Input " << j << ": ";
          cout << hex << setw(8) << setfill('0') << tmp[0][j] << " ";
          cout << hex << setw(8) << setfill('0') << tmp[1][j] << " ";
          cout << hex << setw(8) << setfill('0') << tmp[2][j] << " ";
          cout << hex << setw(8) << setfill('0') << tmp[3][j] << endl;
        }
    }
    
    // 处理剩余未完成的输入
    // for (int j = 0; j < 8; j++) {
    //     if (!finished[j]) {
    //       uint32_t tmp[4][8];
    //       svst1_u32(pg, tmp[0], a_end);
    //       svst1_u32(pg, tmp[1], b_end);
    //       svst1_u32(pg, tmp[2], c_end);
    //       svst1_u32(pg, tmp[3], d_end);
    //       state[j][0] = tmp[0][j];
    //       state[j][1] = tmp[1][j];
    //       state[j][2] = tmp[2][j];
    //       state[j][3] = tmp[3][j];
    //     }
    // }

    // 处理最终结果，转换字节序
    for (int j = 0; j < 8; j++)
    {
        for (int i = 0; i < 4; i++)
        {
            uint32_t value = state[j][i];
            state[j][i] = ((value & 0xff) << 24) |
                          ((value & 0xff00) << 8) |
                          ((value & 0xff0000) >> 8) |
                          ((value & 0xff000000) >> 24);
        }
    }

    // 释放内存
    for (int i = 0; i < 8; i++)
    {
        delete[] paddedMessage[i];
    }
    delete[] messageLength;
}
